package com.zhuodewen.ai.controller.coze;

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.coze.openapi.client.chat.model.ChatEvent;
import com.coze.openapi.client.chat.model.ChatEventType;
import com.coze.openapi.client.connversations.message.model.Message;
import com.coze.openapi.client.files.UploadFileReq;
import com.coze.openapi.client.files.UploadFileResp;
import com.coze.openapi.client.workflows.chat.WorkflowChatReq;
import com.coze.openapi.client.workflows.run.*;
import com.coze.openapi.client.workflows.run.model.WorkflowEvent;
import com.coze.openapi.client.workflows.run.model.WorkflowEventType;
import com.coze.openapi.client.workflows.run.model.WorkflowExecuteStatus;
import com.coze.openapi.client.workflows.run.model.WorkflowRunHistory;
import com.coze.openapi.service.auth.JWTOAuth;
import com.coze.openapi.service.auth.JWTOAuthClient;
import com.coze.openapi.service.config.Consts;
import com.coze.openapi.service.service.CozeAPI;
import com.zhuodewen.ai.base.JSONResult;
import com.zhuodewen.ai.config.coze.CozeConfig;
import com.zhuodewen.ai.config.exception.GlobalException;
import com.zhuodewen.ai.constant.CommonConstants;
import com.zhuodewen.ai.dto.coze.CozeCharflowDto;
import com.zhuodewen.ai.dto.coze.CozeResumeDto;
import com.zhuodewen.ai.dto.coze.CozeUploadDto;
import com.zhuodewen.ai.dto.coze.CozeWorkflowDto;
import com.zhuodewen.ai.vo.coze.CozeWorkflowData;
import com.zhuodewen.ai.vo.coze.WorkflowRunHistoryVo;
import io.reactivex.Flowable;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;

import java.io.InputStream;
import java.io.OutputStream;
import java.net.URL;
import java.net.URLDecoder;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Collections;
import java.util.HashMap;
import java.util.concurrent.TimeUnit;

/**
 * Controller : 控制器类
 * 用于对请求的内容、响应的内容进行数据格式转换。
 */
@RestController                     //等于@ResponseBody(返回JSON格式的数据) + @Controller(定义为Controller接口类)
@RequestMapping(value = "coze")     //路径映射
@Slf4j                              //日志
public class CozeController {

    @Autowired
    private CozeConfig cozeConfig;

    private JWTOAuthClient oauth = null;

    private CozeAPI coze;

    /**
     * 初始化扣子oauth授权
     *
     * @throws Exception
     */
    private void oauth() throws Exception {
        String jwtOauthClientID           = cozeConfig.getCOZE_JWT_OAUTH_CLIENT_ID();
        String jwtOauthPrivateKeyFilePath = cozeConfig.getCOZE_JWT_OAUTH_PRIVATE_KEY_FILE_PATH();
        String jwtOauthPublicKeyID        = cozeConfig.getCOZE_JWT_OAUTH_PUBLIC_KEY_ID();

        //初始化oauth认证和coze操作对象
        String jwtOauthPrivateKey = new String(Files.readAllBytes(Paths.get(jwtOauthPrivateKeyFilePath)), StandardCharsets.UTF_8);
        if (ObjectUtil.isEmpty(oauth) || ObjectUtil.isEmpty(coze)) {
            oauth = new JWTOAuthClient.JWTOAuthBuilder()
                    .clientID(jwtOauthClientID)
                    .privateKey(jwtOauthPrivateKey)
                    .publicKey(jwtOauthPublicKeyID)
                    .baseURL(Consts.COZE_CN_BASE_URL)
                    .build();
            log.info("oauth初始化成功");

            // 官方BUG:https://github.com/coze-dev/coze-java/issues/79, 需要手动指定ttl
            /*coze = new CozeAPI.Builder()
                    .auth(JWTOAuth.builder().jwtClient(oauth).build())
                    .baseURL(Consts.COZE_CN_BASE_URL)
                    .build();*/
            coze = new CozeAPI.Builder()
                    .auth(JWTOAuth.builder().jwtClient(oauth).ttl(oauth.getTtl()).build())
                    .baseURL(Consts.COZE_CN_BASE_URL)
                    .build();
            log.info("coze初始化成功");
        }
    }

    /**
     * 上传文件到扣子
     *
     * @return
     */
    @PostMapping(value = "uploadFileToCoze")
    public JSONResult uploadFileToCoze(@RequestBody CozeUploadDto cozeUploadDto) {
        if (cozeUploadDto.getFile_url().isEmpty()) {
            throw new GlobalException("文件URL不能为空");
        }
        if (cozeUploadDto.getFile_name().isEmpty()) {
            throw new GlobalException("文件名称不能为空");
        }
        // 保存文件到指定目录
        try {
            oauth();
            // 先解码 URL，再修复协议
            String      decodedUrl  = URLDecoder.decode(cozeUploadDto.getFile_url(), StandardCharsets.UTF_8);
            String      fixedUrl    = decodedUrl.replace("http:\\\\", "http://").replace("https:\\\\", "https://");
            InputStream inputStream = new URL(fixedUrl).openConnection().getInputStream();
            // 创建工作流请求(同步/异步)
            UploadFileReq req = UploadFileReq.builder()
                    .fileBytes(inputStream.readAllBytes())
                    .fileName(cozeUploadDto.getFile_name())
                    .build();
            UploadFileResp uploadFileResp = coze.files().upload(req);
            log.info("uploadFileResp:{}", JSONUtil.toJsonStr(uploadFileResp));

            HashMap<String, String> fileMap = new HashMap<>();
            fileMap.put("file_id", uploadFileResp.getFileInfo().getID());
            return new JSONResult().markSuccess(CommonConstants.RESULT_SUCCESS_MSG, fileMap);
        } catch (Exception e) {
            log.error("扣子上传文件API调用失败:{}", e.getMessage());
            throw new GlobalException("扣子上传文件API调用失败:{}" + e.getMessage());
        }
    }

    /**
     * 启动扣子工作流
     * <p>
     * 限制说明:
     * 1)不支持通过此 API 执行包括以下节点的工作流。消息节点、开启了流式输出的结束节点、问答节点
     * 2)通过 API 方式执行工作流前，应确认此工作流已发布，执行从未发布过的工作流时会返回错误码 4200。创建并发布工作流的操作可参考使用工作流。
     * 3)异步运行的参数 is_async 仅限扣子专业版用户使用，否则调用此接口会报错 6003 Workflow execution with is_async=true is a premium feature available only to Coze Professional users.
     * 4)调用此 API 之前，应先在扣子平台中试运行此工作流，如果试运行时需要关联智能体，则调用此 API 执行工作流时，也需要指定智能体ID。通常情况下，执行存在数据库节点、变量节点等节点的工作流需要关联智能体
     *
     * @param cozeWorkflowDto 工作流入参
     */
    @PostMapping(value = "runCozeWorkflow")
    public JSONResult runCozeWorkflow(@RequestBody CozeWorkflowDto cozeWorkflowDto) {
        if (StrUtil.isEmpty(cozeWorkflowDto.getWorkflow_id())) {
            throw new GlobalException("工作流标识workflow_id不能为空");
        }
        try {
            oauth();
            HashMap<String, Object> map = new HashMap<>();
            map.put("input", cozeWorkflowDto.getInput());
            if (StrUtil.isNotEmpty(cozeWorkflowDto.getFile_id())) {
                map.put("file_url", "{\"file_id\": \"" + cozeWorkflowDto.getFile_id() + "\"}");
            }
            // 创建工作流请求(同步/异步)
            RunWorkflowReq req = RunWorkflowReq.builder()
                    .workflowID(cozeWorkflowDto.getWorkflow_id())
                    .parameters(map)
                    .isAsync(cozeWorkflowDto.is_async())
                    .build();
            // 调用应用工作流,需要增加appId
            if (StrUtil.isNotEmpty(cozeWorkflowDto.getApp_id())) {
                req = RunWorkflowReq.builder()
                        .workflowID(cozeWorkflowDto.getWorkflow_id())
                        .parameters(map)
                        .isAsync(cozeWorkflowDto.is_async())
                        .appID(cozeWorkflowDto.getApp_id())
                        .build();
            }
            RunWorkflowResp runWorkflowResp = coze.workflows().runs().create(req);
            log.info("runWorkflowResp:{}", JSONUtil.toJsonStr(runWorkflowResp));
            boolean async = cozeWorkflowDto.is_async();
            if (runWorkflowResp.getCode() != 0) {
                return new JSONResult().markFail(CommonConstants.RESULT_FAIL_MSG, runWorkflowResp.getMsg());
            } else if (runWorkflowResp.getCode() == 0 && async) {
                return new JSONResult().markSuccess(CommonConstants.RESULT_SUCCESS_MSG, new CozeWorkflowData(runWorkflowResp.getExecuteID()));
            } else {
                CozeWorkflowData cozeWorkflowData = JSONUtil.toBean(runWorkflowResp.getData(), CozeWorkflowData.class);
                return new JSONResult().markSuccess(CommonConstants.RESULT_SUCCESS_MSG, cozeWorkflowData);
            }
        } catch (Exception e) {
            log.error("扣子工作流调用失败:{}", e.getMessage());
            throw new GlobalException("扣子工作流调用失败" + e.getMessage());
        }
    }

    /**
     * 查询扣子工作流异步执行结果
     * <p>
     * 限制说明:
     * 1)此接口对于每个工作流的 QPS 限频为 2，即针对某个工作流，每秒钟最多可查看 2 次执行结果。
     * 2)输出节点的输出数据最多保存 24 小时，结束节点为 7 天。
     *
     * @param workflow_id 工作流ID
     * @param execute_id  执行ID
     * @return
     */
    @GetMapping(value = "selectCozeAsyncWorkflow")
    public JSONResult selectCozeAsyncWorkflow(String workflow_id, String execute_id) {
        if (StrUtil.isEmpty(workflow_id)) {
            throw new GlobalException("工作流标识workflow_id不能为空");
        }
        if (StrUtil.isEmpty(execute_id)) {
            throw new GlobalException("执行标识execute_id不能为空");
        }
        boolean              isFinished           = false;
        int                  times                = 0;
        WorkflowRunHistoryVo workflowRunHistoryVo = null;
        WorkflowRunHistory   finalHistory         = null;
        try {
            oauth();
            while (!isFinished && times++ < 10) {
                RetrieveRunHistoryResp historyResp =
                        coze.workflows()
                                .runs()
                                .histories()
                                .retrieve(RetrieveRunHistoryReq.of(workflow_id, execute_id));
                log.info("historyResp:{}", JSONUtil.toJsonStr(historyResp));
                WorkflowRunHistory history = historyResp.getHistories().get(0);
                if (history.getExecuteStatus().equals(WorkflowExecuteStatus.FAIL)) {
                    log.info("Workflow run failed, reason:{},times:{}", history.getErrorMessage(), times);
                    isFinished = true;
                } else if (history.getExecuteStatus().equals(WorkflowExecuteStatus.RUNNING)) {
                    log.info("Workflow run is running,times:{}", times);
                    try {
                        TimeUnit.MILLISECONDS.sleep(500);
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                } else {
                    log.info("Workflow run success:{},times:{}", history.getOutput(), times);
                    isFinished = true;
                }
                if (isFinished || times == 10) {
                    finalHistory = history;
                    break;
                }
            }
            log.info("finalHistory:{}", JSONUtil.toJsonStr(finalHistory));
            JSONObject       outerJson        = JSONUtil.parseObj(finalHistory.getOutput());
            CozeWorkflowData cozeWorkflowData = JSONUtil.toBean(outerJson.getStr("Output"), CozeWorkflowData.class);
            workflowRunHistoryVo = new WorkflowRunHistoryVo(
                    finalHistory.getExecuteID(), finalHistory.getExecuteStatus().toString(), finalHistory.getCreateTime(),
                    finalHistory.getUpdateTime(), cozeWorkflowData, finalHistory.getErrorCode(), finalHistory.getErrorMessage()
            );
        } catch (Exception e) {
            log.error("查询扣子工作流运行状态失败:{}", e.getMessage());
            throw new GlobalException("查询扣子工作流运行状态失败" + e.getMessage());
        }
        return new JSONResult().markSuccess(CommonConstants.RESULT_SUCCESS_MSG, workflowRunHistoryVo);
    }

    /**
     * 启动扣子工作流(流式响应)
     * <p>
     * 限制说明:
     * 1)通过 API 方式执行工作流前，应确认此工作流已发布，执行从未发布过的工作流时会返回错误码 4200。创建并发布工作流的操作可参考使用工作流。
     * 2)调用此 API 之前，应先在扣子平台中试运行此工作流。
     * 如果试运行时需要关联智能体，则调用此 API 执行工作流时，也需要指定 bot_id。通常情况下，执行存在数据库节点、变量节点等节点的工作流需要关联智能体。
     * 执行应用中的工作流时，需要指定 app_id。
     * 请勿同时指定 bot_id 和 app_id，否则 API 会报错 4000，表示请求参数错误。
     * 3)此接口为同步接口，如果工作流整体或某些节点运行超时，智能体可能无法提供符合预期的回复。同步执行时，工作流整体超时时间限制可参考工作流使用限制。
     *
     * @param cozeWorkflowDto 工作流入参
     *                        注意：前端需要使用支持流式响应的HTTP客户端，并逐块读取数据。
     */
    @PostMapping(value = "streamRunCozeWorkflow", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public StreamingResponseBody streamRunCozeWorkflow(@RequestBody CozeWorkflowDto cozeWorkflowDto) {
        if (StrUtil.isEmpty(cozeWorkflowDto.getWorkflow_id())) {
            throw new GlobalException("工作流标识workflow_id不能为空");
        }
        if (StrUtil.isNotEmpty(cozeWorkflowDto.getBot_id()) && StrUtil.isNotEmpty(cozeWorkflowDto.getApp_id())) {
            throw new GlobalException("请勿同时指定bot_id(智能体)和app_id(应用)");
        }
        if (StrUtil.isEmpty(cozeWorkflowDto.getInput())) {
            throw new GlobalException("参数input不能为空");
        }
        HashMap<String, Object> map = new HashMap<>();
        map.put("input", cozeWorkflowDto.getInput());
        if (StrUtil.isNotEmpty(cozeWorkflowDto.getFile_id())) {
            map.put("file_url", "{\"file_id\": \"" + cozeWorkflowDto.getFile_id() + "\"}");
        }
        try {
            oauth();
            RunWorkflowReq req = RunWorkflowReq.builder()
                    .workflowID(cozeWorkflowDto.getWorkflow_id())
                    .readTimeout(300000)    //流式响应设置5分钟超时
                    .parameters(map)
                    .build();
            Flowable<WorkflowEvent> flowable = coze.workflows().runs().stream(req);
            return out -> handleEvent(flowable, coze, out);
        } catch (Exception e) {
            log.error("扣子工作流调用失败:{}", e.getMessage());
            throw new GlobalException("扣子工作流(流式)调用失败" + e.getMessage());
        }
    }

    /**
     * 恢复运行工作流(流式响应)
     * <p>
     * 限制说明:
     * 1)最多调用此接口恢复运行 3 次，如果第三次恢复运行时智能体仍未收到符合预期的回复，则判定为工作流执行失败。
     * 2)恢复运行后，index 和节点 index 都会重置。
     * 3)在扣子专业版中，恢复运行工作流也会产生 token 消耗，且与执行工作流（流式响应）时消耗的 token 数量相同。
     *
     * @param cozeResumeDto
     * @return
     */
    @PostMapping(value = "resumeCozeWorkflow", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public StreamingResponseBody resumeCozeWorkflow(@RequestBody CozeResumeDto cozeResumeDto) {
        if (StrUtil.isEmpty(cozeResumeDto.getWorkflow_id())) {
            throw new GlobalException("工作流标识workflow_id不能为空");
        }
        if (StrUtil.isEmpty(cozeResumeDto.getEvent_id())) {
            throw new GlobalException("中断事件ID event_id不能为空");
        }
        if (ObjectUtil.isEmpty(cozeResumeDto.getInterrupt_type())) {
            throw new GlobalException("中断类型interrupt_type不能为空");
        }
        if (StrUtil.isEmpty(cozeResumeDto.getResume_data())) {
            throw new GlobalException("问题回复resume_data不能为空");
        }
        try {
            oauth();
            ResumeRunReq req = ResumeRunReq.builder()
                    .workflowID(cozeResumeDto.getWorkflow_id())         //工作流ID
                    .eventID(cozeResumeDto.getEvent_id())               //中断事件ID
                    .interruptType(cozeResumeDto.getInterrupt_type())   //中断类型
                    .resumeData(cozeResumeDto.getResume_data())         //问题回复
                    .readTimeout(300000)                                //流式响应设置5分钟超时
                    .build();
            Flowable<WorkflowEvent> resume = coze.workflows().runs().resume(req);
            return out -> handleEvent(resume, coze, out);
        } catch (Exception e) {
            log.error("扣子工作流调用失败:{}", e.getMessage());
            throw new GlobalException("扣子工作流(流式)调用失败" + e.getMessage());
        }
    }

    /**
     * 处理工作流(流式响应)事件
     *
     * @param events
     * @param coze
     * @param out
     */
    private void handleEvent(Flowable<WorkflowEvent> events, CozeAPI coze, OutputStream out) {
        // 不能用官方示例的events.subscribe,必须用blockingForEach,否则会报一些莫名其妙的错误
        events.blockingForEach(
                event -> {
                    try {
                        if (event.getEvent().equals(WorkflowEventType.MESSAGE)) {
                            out.write(event.getMessage().getContent().getBytes());
                            log.info("工作流(流式响应)MESSAGE事件:{}", event.getMessage().getContent());
                        } else if (event.getEvent().equals(WorkflowEventType.ERROR)) {
                            out.write((event.getError().getErrorMessage()).getBytes());
                            log.info("工作流(流式响应)ERROR事件:{}", event.getError().getErrorMessage());
                        } else if (event.getEvent().equals(WorkflowEventType.DONE)) {
                            out.write((ObjectUtil.isEmpty(event.getMessage()) ? StrUtil.EMPTY : event.getMessage().getContent()).getBytes());
                            log.info("工作流(流式响应)MESSAGE事件:{}", ObjectUtil.isEmpty(event.getMessage()) ? StrUtil.EMPTY : event.getMessage().getContent());
                        } else if (event.getEvent().equals(WorkflowEventType.INTERRUPT)) {
                            out.write((ObjectUtil.isEmpty(event.getMessage()) ? StrUtil.EMPTY : event.getMessage().getContent()).getBytes());
                            log.info("工作流(流式响应)MESSAGE事件:{}", ObjectUtil.isEmpty(event.getMessage()) ? StrUtil.EMPTY : event.getMessage().getContent());
                        }
                        out.flush();
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                }
        );
        //coze.shutdownExecutor();
    }

    /**
     * 执行对话流
     * <p>
     * 限制说明:
     * 1)通过 API 方式执行对话流前，应确认此对话流已发布，执行从未发布过的对话流时会返回错误码 4200。如果是扣子应用中的对话流，应先发布扣子应用为 API 服务；如果是空间资源库中的对话流，应先在资源库中发布对话流。
     * 2)此接口暂不支持异步运行。
     *
     * @param cozeCharflowDto
     */
    @PostMapping(value = "streamRunCozeChatflow", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public StreamingResponseBody streamRunCozeChatflow(@RequestBody CozeCharflowDto cozeCharflowDto) {
        if (StrUtil.isEmpty(cozeCharflowDto.getWorkflow_id())) {
            throw new GlobalException("工作流标识workflow_id不能为空");
        }
        if (StrUtil.isNotEmpty(cozeCharflowDto.getBot_id()) && StrUtil.isNotEmpty(cozeCharflowDto.getApp_id())) {
            throw new GlobalException("请勿同时指定bot_id(智能体)和app_id(应用)");
        }
        HashMap<String, Object> map = new HashMap<>();
        map.put("input", cozeCharflowDto.getInput());
        if (StrUtil.isEmpty(cozeCharflowDto.getInput())) {
            throw new GlobalException("参数input不能为空");
        }
        try {
            oauth();
            WorkflowChatReq req = WorkflowChatReq.builder()
                    .workflowID(cozeCharflowDto.getWorkflow_id())
                    .appID(cozeCharflowDto.getApp_id())
                    .botID(cozeCharflowDto.getBot_id())
                    .parameters(map)
                    .additionalMessages(Collections.singletonList(Message.buildUserQuestionText(cozeCharflowDto.getInput())))
                    .readTimeout(300000)    //流式响应设置5分钟超时
                    .build();
            Flowable<ChatEvent> events = coze.workflows().chat().stream(req);
            return out -> handleChatEvent(events, out);
        } catch (Exception e) {
            throw new GlobalException("扣子对话流调用失败" + e.getMessage());
        }
    }

    /**
     * 处理对话流(流式响应)事件
     *
     * @param events
     * @param out
     */
    private void handleChatEvent(Flowable<ChatEvent> events, OutputStream out) {
        events.blockingForEach(
                event -> {
                    if (ChatEventType.CONVERSATION_MESSAGE_DELTA.equals(event.getEvent())) {
                        log.info("event message: {}", event.getMessage().getContent());
                        out.write(event.getMessage().getContent().getBytes());
                        out.flush();
                    }
                });
        //coze.shutdownExecutor();
    }

}
